from datetime import datetime

import backtrader as bt
import akshare as ak
import pandas as pd


"""
创建策略，继承自 backtrader的Strategy
"""
class CAppStrategy(bt.Strategy):
    
    def __init__(self):
        pass


if __name__ == "__main__":
    
    # 初始化回测系统
    cerebro = bt.Cerebro()

    start_date = datetime(2024, 1, 1)
    end_date = datetime(2024, 6, 30)

    print(start_date.strftime("%Y%m%d"))
    print(end_date.strftime("%Y%m%d"))

    # 利用 AkShare 获取后复权数据
    stock_hfq_df = ak.stock_zh_a_daily(symbol="sh601857", start_date=start_date.strftime("%Y%m%d"),
                                       end_date=end_date.strftime("%Y%m%d"), adjust="hfq")
    
    stock_hfq_df.set_index(stock_hfq_df['date'],inplace=True)
    stock_hfq_df.drop('date',axis=1,inplace=True)
    stock_hfq_df.index = pd.to_datetime(stock_hfq_df.index)

    data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)

    # 将数据传入回测系统
    cerebro.adddata(data)

    # 将交易策略加入回测系统
    cerebro.addstrategy(CAppStrategy)

    # 设置初始资本和交易手续费
    start_cash = 10000
    cerebro.broker.setcash(start_cash)
    cerebro.broker.setcommission(commission=0.02)

    # 运行回测系统
    cerebro.run()

    # 获取回测结束后的总资金
    port_value = cerebro.broker.getvalue()  
    pnl = port_value - start_cash  # 盈亏统计
    print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
    print(f"总资金: {round(port_value, 2)}")
    print(f"净收益: {round(pnl, 2)}")
    cerebro.plot(style='candlestick')  # 画图

    #cerebro.plot()